import os
os.environ['DISABLE_ARGUMENTS_CHECKER'] = "Yes"


from nextt.group.group_trades_portfolio import CandidateStockPool, IndexStockPool, MysqlStockPool, GroupStockTrades, update_cfg
from rqalpha.apis import *


def is_sql(sql: str):
    if sql.upper().startswith("SELECT"):
        return True
    return False

def log_cash(context, bar_dict):
    cash_pct = round(context.portfolio.cash / context.init_total_value * 100, 2)
    if cash_pct < 5:
        logger.warn("Remaning cash: %r%%" % cash_pct)
    else:
        logger.info("Remaning cash: %r%%" % cash_pct)
    pass

# 在这个方法中编写任何的初始化逻辑。context对象将会在你的算法策略的任何方法之间做传递。
def init(context):
    from rqalpha.api import scheduler

    context.init_total_value = context.config.base.accounts['STOCK']

    context.per_pos_money = min(context.init_total_value / context.config.params.max_posi_count, context.config.params.per_pos_money)

    start_date = context.config.base.start_date
    end_date = context.config.base.end_date
    context.gt = GroupStockTrades(context.config.params.trades_group_id)
    # 初始化时加载周期内全部交易
    context.gt.load_trades_by_id(start_date=start_date, end_date=end_date)
    # 股池为可选，配置了才加载
    if context.config.params.candidate_pool:
        if isinstance(context.config.params.candidate_pool, str):
            if is_sql(context.config.params.candidate_pool):
                cs = MysqlStockPool(context.config.params.candidate_pool)
                context.candi_func = cs.get_candidate_by_date
            elif context.config.params.candidate_pool.endswith('.XSHE') or context.config.params.candidate_pool.endswith('.XSHG'):
                cs = IndexStockPool(context.config.params.candidate_pool, start_date=start_date, end_date=end_date)
                context.candi_func = cs.get_candidate_by_date
            else:
                cs = CandidateStockPool(context.config.params.candidate_pool, start_date=start_date, end_date=end_date)
                context.candi_func = cs.get_candidate_by_date
        else:
            context.candi_func = context.config.params.candidate_pool
    else:
        context.candi_func = None
    # 每个月一次观察仓位
    scheduler.run_monthly(log_cash, tradingday=1)
    pass

# before_trading此函数会在每天交易开始前被调用，当天只会被调用一次
def before_trading(context):
    # 打印当前时间日志
    context.trades_open = context.gt.get_open_trades_by_date(context.now.date())
    context.trades_close = context.gt.get_close_trades_by_date(context.now.date())

# 你选择的期货数据更新将会触发此段逻辑，例如日线或分钟线更新
def handle_bar(context, bar_dict):
    # 当日无交易直接返回
    if context.trades_open is None and context.trades_close is None:
        return
    # 目前股票暂支持日线
    if context.config.base.frequency == '1d':
        # 股池为可选，配置了才加载
        if context.candi_func:
            candidate_df = context.candi_func(context.now.date())
            candidates = candidate_df['code'].tolist()
        else:
            candidates = None
        # 先卖出不在候选股池票
        posi_codes = []
        positions = get_positions()
        # 如果有股池，不在股池先卖出
        if candidates is not None:
            assert(len(candidates) > 0)
            for p in positions:
                if p.order_book_id not in candidates:
                    order_target_percent(p.order_book_id, 0)
                else:
                    posi_codes.append(p.order_book_id)
        else:
            posi_codes = [p.order_book_id for p in positions]
        # 再卖出卖出信号票     
        if context.trades_close is not None:
            for idx,row in context.trades_close.iterrows():
                if row.order_book_id in posi_codes:
                    order_target_percent(row.order_book_id, 0, price=row.last_price)
        # 重新配置股池仓位，只处理开仓票，不动态再平衡，减少换手率
        if context.trades_open is not None:
            # 有候选股池只买股池票
            if candidates is not None:
                context.trades_open = context.trades_open[context.trades_open['order_book_id'].isin(candidates)]
            open_codes = context.trades_open['order_book_id'].tolist()
            total_codes = set(posi_codes + open_codes)
            if len(total_codes) > context.config.params.max_posi_count:
                if context.config.params.over_count_proc == 1:
                    if candidates is not None:
                        # drop low score, 超出最大，去掉低分
                        candi_all = [c for c in candidates if c in total_codes]
                        # 默认低分值排后面
                        candi_high = candi_all[0:context.config.params.max_posi_count]
                        for c in posi_codes:
                            if c not in candi_high:
                                order_target_percent(c, 0)  # 无卖出价格，只能收盘价卖出
                        # 复利按比例，否则定额
                        candi_open = context.trades_open[context.trades_open['order_book_id'].isin(candi_high)]
                        if context.config.params.is_compound:
                            posi_percent = 1 / context.config.params.max_posi_count
                            for idx,row in candi_open.iterrows():
                                if row.order_book_id not in posi_codes:
                                    order_target_percent(row.order_book_id, posi_percent, price=row.last_price)
                        else:
                            for idx,row in candi_open.iterrows():
                                if row.order_book_id not in posi_codes:
                                    order_target_percent(row.order_book_id, context.per_pos_money, price=row.last_price)
                    else:
                        if context.config.params.is_compound:
                            posi_percent = 1 / context.config.params.max_posi_count
                            for idx,row in context.trades_open.iterrows():
                                if row.order_book_id not in posi_codes:
                                    order_target_percent(row.order_book_id, posi_percent, price=row.last_price)
                        else:
                            # 定额按顺序买入，收益存在偶然性
                            for idx,row in context.trades_open.iterrows():
                                if row.order_book_id not in posi_codes:
                                    order_target_percent(row.order_book_id, context.per_pos_money, price=row.last_price)
                elif context.config.params.over_count_proc == 2:
                    # equal weight, 等权满仓
                    if context.config.params.is_compound:
                        posi_percent = min(1 / context.config.params.min_posi_count, 1 / len(total_codes))
                        for idx,row in context.trades_open.iterrows():
                            if row.order_book_id not in posi_codes:
                                order_target_percent(row.order_book_id, posi_percent, price=row.last_price)
                    else:
                        # 定额按顺序买入，收益存在偶然性
                        for idx,row in context.trades_open.iterrows():
                            if row.order_book_id not in posi_codes:
                                order_target_value(row.order_book_id, context.per_pos_money, price=row.last_price)
                else:
                    # cash equal weight, 可用现金等权满仓
                    if context.config.params.is_compound:
                        open_codes = [o for o in open_codes if o not in posi_codes]
                        posi_percent = min(1 / context.config.params.min_posi_count, 1 / len(total_codes), context.portfolio.cash / context.portfolio.total_value / len(open_codes))
                        for idx,row in context.trades_open.iterrows():
                            if row.order_book_id not in posi_codes:
                                ret = order_target_percent(row.order_book_id, posi_percent, price=row.last_price)
                                if ret is None:
                                    print(context.portfolio.cash, context.portfolio.total_value, posi_percent)
                    else:
                        # 定额按顺序买入，收益存在偶然性
                        for idx,row in context.trades_open.iterrows():
                            if row.order_book_id not in posi_codes:
                                order_target_value(row.order_book_id, context.per_pos_money, price=row.last_price)
                    pass
            elif len(open_codes) > 0:
                if context.config.params.is_compound:
                    # 复利考虑资金分配比例
                    if context.config.params.less_count_proc == 1:
                        # lower position， 按配置仓位，只买新进股票
                        posi_percent = 1 / context.config.params.max_posi_count
                        for idx,row in context.trades_open.iterrows():
                            if row.order_book_id not in posi_codes:
                                order_target_percent(row.order_book_id, posi_percent, price=row.last_price)
                    elif context.config.params.less_count_proc == 2:
                        # equal weight, 等权满仓
                        posi_percent = min(1 / context.config.params.min_posi_count, 1 / len(total_codes))
                        for idx,row in context.trades_open.iterrows():
                            if row.order_book_id not in posi_codes:
                                order_target_percent(row.order_book_id, posi_percent, price=row.last_price)
                    else:
                        # cash equal weight, 可用现金等权满仓
                        open_codes = [o for o in open_codes if o not in posi_codes]
                        posi_percent = min(1 / context.config.params.min_posi_count, 1 / len(total_codes), context.portfolio.cash / context.portfolio.total_value / len(open_codes))
                        for idx,row in context.trades_open.iterrows():
                            if row.order_book_id not in posi_codes:
                                ret = order_target_percent(row.order_book_id, posi_percent, price=row.last_price)
                                if ret is None:
                                    print(context.portfolio.cash, context.portfolio.total_value, posi_percent)
                        pass
                else:
                    for idx,row in context.trades_open.iterrows():
                        if row.order_book_id not in posi_codes:
                            order_target_value(row.order_book_id, context.per_pos_money, price=row.last_price)
    else:
        pass
        # # 股票分钟信号买卖
        # if context.trades_close is not None:
        #     trades = context.trades_close.loc[context.now]
        #     if isinstance(trades, pd.Series):
        #         order_target_percent(trades.order_book_id, 0)
        #     else:
        #         for idx,row in trades.iterrows():
        #             order_target_percent(row.order_book_id, 0)
        # #
        # if context.trades_open is not None:
        #     trades = context.trades_open.loc[context.now]
        #     if isinstance(trades, pd.Series):
        #         order_target_percent(trades.order_book_id, 0.01, price=trades.last_price)
        #     else:
        #         for idx,row in trades.iterrows():
        #             order_target_percent(row.order_book_id, 0.01, price=row.last_price)



__config__ = {
  "base": {
    "start_date": "2012-01-01",
    "end_date": "2021-12-31",
    "frequency": "1d",
    "forced_liquidation": False,
    "accounts": {
        "stock": 200000_0000
    }
  },
  "params": {
    "is_compound": True,   # 是否复利，复利按比例，非复利即定额
    "trades_group_id": '19ab21a4-a2ae-4d07-a94a-000cfefc5797',  # 交易组id
    "candidate_pool": "select day,code,market_cap from tb_jq_fundamental_valuation where market_cap < 80.0;",  # 股池名或函数，None表示全市场，不用股池过滤
    "max_posi_count": 2000, # 最大持仓数量
    "min_posi_count": 500, # 最小持仓数量
    "per_pos_money": 200_0000, # 初始每次下单额
    "over_count_proc": 3,  # 1.drop low score, 2.equal weight, 3.cash equal weight
    "less_count_proc": 3, # 1.lower position, 2.equal weight, 3.cash equal weight
  },
  "extra": {
    "log_level": "warning"
  },
  "mod": {
    "sys_analyser": {
      "benchmark": "000906.XSHG",
      "enabled": True,
      "plot": True
    },
    "sys_simulation": {
      "enabled": True,
      "signal": True,
      "price_limit": True
    },
    "stu_log": {
      "enabled": False,
      "ip":"192.168.12.110"
    }
  }
}


def run_gst(param={}):
    # 可让外部调用回测函数
    cfg = update_cfg(__config__, param)
    #
    start = time.time()
    print(cfg)
    ret = run_func(init=init, before_trading=before_trading, handle_bar=handle_bar, config=cfg)
    print(ret)
    print(time.time() - start)
    return ret


if __name__ == '__main__':
    import time
    import json

    from rqalpha import run_func
    # 您可以指定您要传递的参数
    # import cProfile
    # pr = cProfile.Profile()
    # pr.enable()
    ret = run_gst()
    pd.set_option('display.max_rows', None)
    trades = ret['sys_analyser']['trades']
    print(trades)
    print(len(trades), trades['transaction_cost'].sum(), trades['transaction_cost'].mean(), trades['transaction_cost'].min(), trades['transaction_cost'].max())
    # pr.disable()
    # pr.print_stats('cumtime')

    # from nextt.tasks import year_split_file_tasks
    
    # param_list = [
    #     {"max_posi_count": 100, "over_count_proc":1, "less_count_proc": 1},
    #     {"max_posi_count": 100, "over_count_proc":1, "less_count_proc": 2},
    #     {"max_posi_count": 100, "over_count_proc":2, "less_count_proc": 1},
    #     {"max_posi_count": 100, "over_count_proc":2, "less_count_proc": 2}
    #     ] 
    # start = time.time()
    # ret = year_split_file_tasks(__file__, param_list, __config__, 40)
    # print(time.time() - start)
    # print(json.dumps(ret, indent=2))